/* 
 * This code is part of project 
 * "Sårbarheder i og trusler mod trådløse netværk med fokus på WEP."
 * Projects in IT-Security and Cryptography - Department of Computer Science - Aarhus University 
 * April - June 2015
 * 
 * Jesper René Frandsen 
 * George Catalin Hategan 
 * Torben Krogh Johansen
 */
package krypto.projekt;

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Random;

/**
 * This is a implementation of a RC4 attack
 * 
 * The overall structure is this:
 * - Generate a random key (int[] key) which we will try to guess using the FMS attack.
 * - Outer loop (see integer allBytesinKeyLoop) - controlling which key value is to be guessed now.
 * --- Inner loop (see integer oneByteInKeyLoop) - the actual guessing of a key value.
 * 
 * The FMS attack describe a RC4 internal state that makes it possible to guess the key.
 * This state is called resolved state.
 * 
 * When the state is resolved, we store the guess in a hashmap. 
 * At the end of the inner loop, we decide which guess has been the most popular (most guessed at value).
 * 
 * At the end a final status is printed on stdout.
 * It contains the key we try to guess. The actual guess. Time it took to guess.
 */
public class RC4Crack {
	public static void main(String[] args) throws Exception {
		final boolean DEBUG = false;

		long stopwatchStart, stopwatchStop;
		stopwatchStart = System.currentTimeMillis();

		final int KEY_LENGTH = 10;
		int[] key = new int[KEY_LENGTH];
		int[] keyGuess = new int[KEY_LENGTH];

		final int KNOWN_KEY_LENGTH = 3;
		int numberOfKeysToGuess = KEY_LENGTH - KNOWN_KEY_LENGTH;

		final int LOOP_NUMBER = 10000000;

		//making a random key array to guess and initialize output array to zeros.
		for (int i=0; i<KEY_LENGTH; i++) {
			key[i] = getRandomBetween0and255();
			keyGuess[i] = 0;
		}

		int[] keyForResolvedTest;

		System.out.println("RC4Crack started - LOOP_NUMBER=" + LOOP_NUMBER);

		int guessingKeyIndex = KNOWN_KEY_LENGTH;

		for (int allBytesinKeyLoop = 0; allBytesinKeyLoop < numberOfKeysToGuess; allBytesinKeyLoop++) {
			Map<Integer, Integer> map = new HashMap<Integer, Integer>();

			int isResolvedCount = 0;
			int isNotResolvedCount = 0;
			int isBingoCount = 0;

			for (int oneByteInKeyLoop = 0; oneByteInKeyLoop < LOOP_NUMBER; oneByteInKeyLoop++) {
				key[0] = getRandomBetween0and255(); //Typical use of IV in WEP
				key[1] = getRandomBetween0and255(); //Typical use of IV in WEP
				key[2] = getRandomBetween0and255(); //Typical use of IV in WEP

				//Generate a stream to get element x0, this is the only place where we actively use the real key.
				RC4Impl RC4S = new RC4Impl();	
				RC4S.doKSA(key);
				int x0 = RC4S.getNextFromPRGA();

				//Check to see if the internal state is resolved
				keyForResolvedTest = new int[guessingKeyIndex];
				for (int i=0; i<guessingKeyIndex; i++) {
					//only IV's taken from key, rest is taken from guessed values 
					if(i<3) {
						keyForResolvedTest[i] = key[i];
					} else {
						keyForResolvedTest[i] = keyGuess[i];
					}
				}

				RC4Impl RC4 = new RC4Impl();	

				if (RC4.isResolved(keyForResolvedTest, guessingKeyIndex , x0) == true) {
					isResolvedCount++;
					if (DEBUG) {
						System.out.print("iv+key:");

						for (int i=0; i<key.length; i++ ) {
							System.out.print(" " + key[i]);
						}

						System.out.print(",x0=" + x0 + ",position=" + guessingKeyIndex + ",KeyGuess=" + RC4.getKeyGuess());
					}

					//If the guess is correct, we increase correct guess counter. This is just for statistics.
					if (RC4.getKeyGuess() == key[guessingKeyIndex]) {
						if (DEBUG)
							System.out.print(" - Bingo!");

						isBingoCount++;
					}

					if (DEBUG)
						System.out.println("");

					//Add value to hashmap with increased count value
					addValueToMap(map, RC4.getKeyGuess());
				} else {
					isNotResolvedCount++;
				}
			}

			//Print status
			System.out.println("guessingKeyIndex=" + guessingKeyIndex + ",isNotResolvedCount=" + 
					isNotResolvedCount  + ",isResolvedCount=" + isResolvedCount + ",isBingoCount=" + isBingoCount);

			//Find the value guessed at most time
			int mostPopular = getMostPopularValueFromMap(map);

			System.out.println("Most popular=" + mostPopular);
			keyGuess[guessingKeyIndex] = mostPopular;

			//Indicate we are finished guessing this byte of the key, and move forward to the next byte.
			guessingKeyIndex++;
		}

		//Print final status
		System.out.print("Key to guess is :");		
		for (int i=KNOWN_KEY_LENGTH; i<key.length && i<numberOfKeysToGuess+KNOWN_KEY_LENGTH; i++ ) {
			System.out.print(" " + key[i]);
		}		
		System.out.println("");

		System.out.print("Guessed key is  :");
		for (int i=KNOWN_KEY_LENGTH; i<keyGuess.length && i<numberOfKeysToGuess+KNOWN_KEY_LENGTH; i++ ) {
			System.out.print(" " + keyGuess[i]);
		}
		System.out.println("");

		stopwatchStop = System.currentTimeMillis();
		System.out.println("Operation lasted " + (stopwatchStop - stopwatchStart)/1000 + " seconds.");
	}

	/**
	 * Adding value and value count to hashmap.
	 * If value is present, value count is increased.
	 * If value is not present, it is added with a value count of 1.
	 */
	private static void addValueToMap(Map<Integer, Integer> map, int  valueToAdd) {
		Integer count = map.get(valueToAdd);
		if (count == null) {
			count = 1;
		} else {
			count++;
		}

		map.put(valueToAdd , count);
	}

	/**
	 * Getting the most popular value from a hashmap.
	 * If the hashmap contains to equally popular values, the first one is returned.
	 */
	private static int getMostPopularValueFromMap(Map<Integer, Integer> map) {
		// make sure map contains at least one value
		Integer count = map.get(0);
		if (count == null) {
			map.put(0, 0);
		} 

		// Return max value in the hashmap
		int maxValueInMap=(Collections.max(map.values()));  

		int mostPopular = 0;

		// Iterate through hashmap to find entry matching value
		for (Entry<Integer, Integer> entry : map.entrySet()) {  
			if (entry.getValue()==maxValueInMap) {
				mostPopular= entry.getKey();     
			}
		}

		return mostPopular;
	}

	/**
	 * Get a random value between 0 and 255
	 */
	private static int getRandomBetween0and255() {
		Random random = new Random();

		int min = 0;
		int max = 255;

		return random.nextInt(max - min + 1) + min;
	}
}
